from torch import nn
from torch.autograd import Function
from torch.autograd.function import once_differentiable

from taa_core._C import pred_rand_gt


class RandMatchGT(Function):
    @staticmethod
    def forward(ctx, pred_gt_loss, non_label, ratio):
        return pred_rand_gt(pred_gt_loss, non_label, ratio)

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_outputs):
        return None


rand_match_gt = RandMatchGT.apply


class AnchorMatcher(nn.Module):
    def __init__(self):
        super(AnchorMatcher, self).__init__()

    def forward(self, pred_gt_loss, non_label, ratio=0.0):
        return rand_match_gt(pred_gt_loss, non_label, ratio)